# import math
import numpy as np
import pandas as pd


class DataLoader:
    """A class for loading and transforming data for the lstm model"""

    def __init__(self, configs):
        price_column = configs['data']['price']['column']
        split = configs['data']['train_test_split']
        window = configs['predict']['window']
        sampling_frequency = configs['predict']['sampling_frequency']

        dataframe = pd.read_csv(configs['data']['price']['filename'])
        y_column = ''
        if configs['predict']['signal'] == 'price':
            y_column = price_column
        elif configs['predict']['signal'] == 'return':
            dataframe['return'] = (dataframe[price_column].shift(-window) - dataframe[price_column].shift(-1)) \
                                  / dataframe[price_column].shift(-1)
            y_column = 'return'
        elif configs['predict']['signal'] == 'volatility':
            dataframe['volatility'] = dataframe[price_column].rolling(window).std().shift(-window)
            y_column = 'volatility'
        else:
            print('Unsupported predictive signal type, please check!')
            exit()
        dataframe = dataframe[::sampling_frequency]
        dataframe = dataframe[:-1]
        dataframe = dataframe.get(['time', y_column])
        for path in configs['data']['features_list']:
            data = pd.read_csv(path)
            data = data[::sampling_frequency]
            dataframe = dataframe.merge(data, on='time', how='left')
        if ~isinstance(split, float):
            i_split = dataframe.loc[dataframe['time'] == split].index[0]
        else:
            i_split = int(len(dataframe) * split)
        cols = list(dataframe.columns.values)
        cols.remove('time')
        self.data_train = dataframe.get(cols).values[:i_split]
        self.data_test = dataframe.get(cols).values[i_split:]
        self.len_train = len(self.data_train)
        self.len_test = len(self.data_test)
        self.len_train_windows = None
        seq_len = configs['data']['sequence_length']
        self.test_time = dataframe.get(['time', y_column]).values[i_split + (seq_len - 1):]
        self.test_y0 = dataframe.get(y_column).values[i_split:-seq_len + 1]

    def get_test_time(self):
        return self.test_time

    def get_test_y0(self):
        return self.test_y0

    def get_test_data(self, seq_len, normalise):
        """
        Create x, y test data windows
        Warning: batch method, not generative, make sure you have enough memory to
        load data, otherwise reduce size of the training split.
        """
        data_windows = []
        for i in range(self.len_test - seq_len + 1):
            data_windows.append(self.data_test[i:i + seq_len])

        data_windows = np.array(data_windows).astype(float)
        data_windows = self.normalise_windows(data_windows, single_window=False) if normalise else data_windows

        x = data_windows[:, :-1]
        y = data_windows[:, -1, [0]]
        return x, y

    def get_train_data(self, seq_len, normalise):
        """
        Create x, y train data windows
        Warning: batch method, not generative, make sure you have enough memory to
        load data, otherwise use generate_training_window() method.
        """
        data_x = []
        data_y = []
        for i in range(self.len_train - seq_len):
            x, y = self._next_window(i, seq_len, normalise)
            data_x.append(x)
            data_y.append(y)
        return np.array(data_x), np.array(data_y)

    def generate_train_batch(self, seq_len, batch_size, normalise):
        """Yield a generator of training data from filename on given list of cols split for train/test"""
        i = 0
        while i < (self.len_train - seq_len):
            x_batch = []
            y_batch = []
            for b in range(batch_size):
                if i >= (self.len_train - seq_len):
                    # stop-condition for a smaller final batch if data doesn't divide evenly
                    yield np.array(x_batch), np.array(y_batch)
                    i = 0
                x, y = self._next_window(i, seq_len, normalise)
                x_batch.append(x)
                y_batch.append(y)
                i += 1
            yield np.array(x_batch), np.array(y_batch)

    def _next_window(self, i, seq_len, normalise):
        """Generates the next data window from the given index location i"""
        window = self.data_train[i:i + seq_len]
        window = self.normalise_windows(window, single_window=True)[0] if normalise else window
        x = window[:-1]
        y = window[-1, [0]]
        return x, y

    @staticmethod
    def normalise_windows(window_data, single_window=False):
        """Normalise window with a base value of zero"""
        normalised_data = []
        window_data = [window_data] if single_window else window_data
        for window in window_data:
            normalised_window = []
            for col_i in range(window.shape[1]):
                index = 0
                num = float(window[index, col_i])
                while num == 0:
                    index += 1
                    num = float(window[index, col_i])
                normalised_col = [((float(p) / num) - 1) for p in window[:, col_i]]
                normalised_window.append(normalised_col)
            normalised_window = np.array(
                normalised_window).T  # reshape and transpose array back into original multidimensional format
            normalised_data.append(normalised_window)
        return np.array(normalised_data)
